#pragma once
#include<Eigen/Dense>
#include<opencv2/core/eigen.hpp>
#include"utils/yolo.h"
#include"utils/yolo_utils.h"
class YOLOv8Seg : public yolo::YOLO
{
public:
	YOLOv8Seg(const yolo_utils::InitParameter& param);
	~YOLOv8Seg();
	virtual bool init(const std::vector<unsigned char>& trtFile);
	virtual void preprocess(const std::vector<cv::Mat>& imgsBatch);
	virtual bool infer();
	virtual void postprocess(const std::vector<cv::Mat>& imgsBatch);
	virtual void reset();

public:
	void showAndSave(const std::vector<std::string>& classNames,
		const int& cvDelayTime, std::vector<cv::Mat>& imgsBatch);
    void setmask(const int& cvDelayTime, cv::Mat &img, cv::Mat &mask);
    void detect_target(cv::Mat &mask, const cv::Mat &img);

private:
	float* m_output_src_transpose_device;
	float* m_output_seg_device; // eg:116 * 8400, 116=4+80+32
	float* m_output_objects_device;

	float* m_output_seg_host;
	float* m_output_objects_host;

	int m_output_objects_width; // 39 = 32 + 7, 7:left, top, right, bottom, confidence, class, keepflag; 
	int m_output_src_width; // 116 = 4+80+32, 4:xyxy; 80:coco label; 32:seg
	nvinfer1::Dims m_output_seg_dims;
	int m_output_obj_area;
	int m_output_seg_area;
	int m_output_seg_w;
	int m_output_seg_h;

	cv::Mat m_mask160; 
	Eigen::MatrixXf m_mask_eigen160;
	cv::Rect m_thresh_roi160;
	cv::Rect m_thresh_roisrc;
	float m_downsample_scale;
	cv::Mat m_mask_src;
	cv::Mat m_img_canvas;

    std::vector<int> target_label_coco80 = {0, 1, 2, 3, 5, 7};  // for coco80
    // 定义腐蚀的核（structuring element）
    cv::Mat kernel = cv::getStructuringElement(cv::MORPH_RECT, cv::Size(10, 10));
};